{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "client = OpenAI(\n",
    "    api_key=\"ollama\",\n",
    "    base_url=\"http://192.168.20.43:11434/v1\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"管道名称\": \"长输管道黄河段\",\n",
      "  \"管道描述\": \"全长2.7km，材质为L245NB钢。\",\n",
      "  \"管道口径\": \"189×5mm\",\n",
      "  \"管道压力\": \"None\"\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "def llm_extract_info(query, info_escription, examples):\n",
    "    prompt = f\"\"\"\n",
    "    ### 目标\n",
    "    从用户的query中提取出关键信息。并按Example中的格式输出，仅输出<output></output>中的内容。\n",
    "\n",
    "    ### 什么是关键信息\n",
    "    {info_escription}\n",
    "\n",
    "    ### Examples\n",
    "    {examples}\n",
    "\n",
    "    ### 约束\n",
    "    - 不要编造事实。如果用户的query中没有提及相关信息，则设为None。\n",
    "    \"\"\"\n",
    "\n",
    "    query_new = f\"\"\"\n",
    "    <query>{query}</query>\n",
    "    \"\"\"\n",
    "    \n",
    "    completion = client.chat.completions.create(\n",
    "        model=\"qwen2.5:14b\",\n",
    "        messages=[{'role': 'system', 'content': prompt},\n",
    "                  {'role': 'user', 'content': query_new}],\n",
    "        )\n",
    "    return completion.choices[0].message.content\n",
    "\n",
    "\n",
    "# 测试用例\n",
    "info_description = \"\"\"\n",
    "管道名称，参数类型String，管道的编号或名称。例如：新疆LST-0102-A\n",
    "管道描述，参数类型String，所有query中与该管道相关的描述信息。例如：某出油管道于2013年投运，全长0.7km，管道材质为L245NB钢。\n",
    "管道口径，参数类型NUMBER，指的是管道内口径，单位是mm。例如：89×5mm\n",
    "管道压力，参数类型NUMBER，指的是实际运行压力而非设计时压力。单位是MPa。例如：6.4MPa\n",
    "\"\"\"\n",
    "info_examples = \"\"\"\n",
    "<example>\n",
    "<query>新疆LST-0102-A管道，2013年投运，全长0.7km，材质L245NB钢，89×5mm口径，6.4MPa压力。</query>\n",
    "<output>\n",
    "{\n",
    "  \"管道名称\": \"新疆LST-0102-A\",\n",
    "  \"管道描述\": \"某出油管道于2013年投运，全长0.7km，管道材质为L245NB钢。\",\n",
    "  \"管道口径\": \"89×5mm\",\n",
    "  \"管道压力\": \"6.4MPa\"\n",
    "}\n",
    "</output>\n",
    "</example>\n",
    "<example>\n",
    "<query>请问LST-0102-A的管道口径是多少？</query>\n",
    "<output>\n",
    "{\n",
    "  \"管道名称\": \"LST-0102-A\",\n",
    "  \"管道描述\": \"None\",\n",
    "  \"管道口径\": \"None\",\n",
    "  \"管道压力\": \"None\"\n",
    "}\n",
    "</output>\n",
    "</example>\n",
    "<example>\n",
    "<query>某出油管道于2013年投运，全长0.7km，管道材质为L245NB钢，管线规格为Φ89×5mm，设计压力6.4MPa，最大输送量10万吨/年，近三年最大运行压力0.7MPa，外防腐层材料为3PE，采用牺牲阳极保护系统。输送介质为原油，平均含水量43%（体积分数），历史上发生过4次内腐蚀失效，既往未开展过任何检测评价工作。</query>\n",
    "<output>\n",
    "{\n",
    "  \"管道名称\": \"None\",\n",
    "  \"管道描述\": \"某出油管道于2013年投运，全长0.7km，管道材质为L245NB钢，最大输送量10万吨/年，外防腐层材料为3PE，采用牺牲阳极保护系统。输送介质为原油，平均含水量43%（体积分数），历史上发生过4次内腐蚀失效，既往未开展过任何检测评价工作。\",\n",
    "  \"管道口径\": \"89×5mm\",\n",
    "  \"管道压力\": \"0.7MPa\"\n",
    "}\n",
    "</output>\n",
    "</example>\n",
    "\"\"\"\n",
    "\n",
    "query1 = \"\"\"\n",
    "某集油管道为例，该管道于1993年投运，全长100km，管道材质为L415钢，管线规格为Φ426×10mm，设计压力6.4MPa，最大输送量300万吨/年，近三年最大运行压力曾达到6.3MPa，外防腐层材料为石油沥青，采用强制电流阴极保护。输送介质为含硫0.873%，含盐47.78mg/L，含水1.06%的原油。\n",
    "\"\"\"   \n",
    "query2 = \"\"\"\n",
    "LST-0102-A管道，运行压力6.4MPa，需要检查。\n",
    "\"\"\"   \n",
    "query3 = \"\"\"\n",
    "长输管道黄河段，全长2.7km，材质L245NB钢，189×5mm口径，设计压力1.5MPa。\n",
    "\"\"\"   \n",
    "\n",
    "info = llm_extract_info(query3, info_description, info_examples) \n",
    "print(info)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
